import numpy as np

def median_filter_1d(data, kernel_size):
    '''
    中值滤波
    :param data:    数据源
    :param kernel_size:     滤波器长度
    :return:
    '''
    window = np.array(data[-kernel_size-1:-1])
    return np.median(window,axis=0)

class Kalman_Filter3D:
    '''
        三维，无方向场景下的卡尔曼滤波算法模组
    '''
    def __init__(self):
        # 初始状态                 x  y  z vx vy vz
        self.initState=np.array([0, 0, 0, 0, 0, 0],dtype=np.float)
        # 初始协方差,可以看出是每个维度都是一一对应的关系
        '''
        [ 1 0 0 0 0 0 ]
        [ 0 1 0 0 0 0 ]
        [ 0 0 1 0 0 0 ]
        [ 0 0 0 1 0 0 ]
        [ 0 0 0 0 1 0 ]
        [ 0 0 0 0 0 1 ]
        '''
        self.initCov=np.eye(6)
        #   状态转移矩阵
        self.stateTransMatrix=np.array([[1,0,0,1,0,0],
                                      [0,1,0,0,1,0],
                                      [0,0,1,0,0,1],
                                      [0,0,0,1,0,0],
                                      [0,0,0,0,1,0],
                                      [0,0,0,0,0,1]],dtype=float)
        #   观测矩阵                    X  Y  Z  Vx Vy Vz
        self.observeMatrix=np.array([[1, 0, 0, 0, 0, 0],
                                     [0, 1, 0, 0, 0, 0],
                                     [0, 0, 1, 0, 0, 0]],dtype=float)
        #   过程噪声（先设定一个初始值，这个需要跟据你系统的评估来确定）

        self.InitParams()
    def InitParams(self):
        '''
        初始化状态变量
        :return:
        '''
        self.currentState=self.initState.copy()
        self.predictState=self.initState.copy()
        self.currentCov=self.initCov
        self.predictedCov=self.currentCov

    def Predict(self,velocity=np.array([0,0,0],dtype=float),procNoise=0.02):
        '''
            预测过程
            :param v:
            :return:
        '''
        #   把速度赋值给状态中的“速度”属性
        self.currentState[3:6] = velocity
        #   基于当前的速度，预测下一个状态的状态数值
        self.predictState=self.stateTransMatrix.dot(self.currentState)
        #   预测三维环境下的协方差矩阵
        self.predictedCov=self.stateTransMatrix.dot(self.currentCov).dot(self.stateTransMatrix.T)+np.eye(6)*procNoise

    def Update(self,observed_Pos=np.array([0,0,0],dtype=float),observeNoise=1.0):
        '''
        更新数据
            :param observed_Pos: 带有误差的位置观测值
            :return:
        '''
        #   卡尔曼增益（Kalman Gain）计算
        '''
            K=\frac{估计的误差}{估计的误差+测量的误差}=\frac{\hat{P_k}C}{C\hat{P_k}C^T+Error}
        '''
        self.Kalman_Gain = self.predictedCov.dot(self.observeMatrix.T) \
            .dot(np.linalg.inv( \
            self.observeMatrix.dot(self.predictedCov).dot(self.observeMatrix.T) + np.eye(3)*observeNoise))

        '''
            基于Kalman Gain估算当前状态
        '''
        self.currentState = self.predictState + self.Kalman_Gain.dot(observed_Pos-self.observeMatrix.dot(self.predictState))

        '''
            当前协方差估计
        '''
        self.currentCov = (np.eye(6) - self.Kalman_Gain.dot(self.observeMatrix)).dot(self.predictedCov)

        return self.currentState[0:3]